#!/usr/bin/env python3
# Author: Armit
# Create Time: 周一 2025/07/21

# 观察数据集统计特征

import json
from pathlib import Path
from collections import defaultdict
from pprint import pprint as pp

BASE_PATH = Path(__file__).parent
DATA_PATH = BASE_PATH / 'data'
OUT_PATH = BASE_PATH / 'out'
OUT_PATH.mkdir(exist_ok=True)

PLT_NAME = 'plant'
CLS_NAME = 'class'
ENV_NAME = 'environment'
ARE_NAME = 'area'
ALT_NAME = 'altitude'


def stats_trainset():
  PHRASES = {
    PLT_NAME: set(),
    CLS_NAME: set(),
    ENV_NAME: set(),
    ARE_NAME: set(),
    ALT_NAME: set(),
  }
  for i in range(1, 5+1):
    fp = DATA_PATH / 'train' / f'train0{i}.jsonl'
    items = []
    with open(fp, 'r', encoding='utf-8') as fh:
      for line in fh.readlines():
        it = json.loads(line.strip())
        items.append(it)
    print(f'>> file {fp}')
    print(f'>> nlen: {len(items)}')

    cntr = {}
    for it in items:
      for key, val in it['label'].items():
        if key not in cntr:
          cntr[key] = defaultdict(int)
        for name, lcut in val.items():
          PHRASES[key].add(name)
          cntr[key][name] += 1

    save_fp = OUT_PATH / f'{fp.stem}-phrases.json'
    print(f'>> save file: {save_fp}')
    with open(save_fp, 'w', encoding='utf-8') as fh:
      json.dump(cntr, fh, indent=2, ensure_ascii=False)

  save_fp = OUT_PATH / 'trainset-phrases.json'
  print(f'>> save file: {save_fp}')
  with open(save_fp, 'w', encoding='utf-8') as fh:
    json.dump({k: sorted(v) for k, v in PHRASES.items()}, fh, indent=2, ensure_ascii=False)


def stats_testset():
  fp = DATA_PATH / 'test.jsonl'
  items = []
  with open(fp, 'r', encoding='utf-8') as fh:
    for line in fh.readlines():
      it = json.loads(line.strip())
      items.append(it)
  print(f'>> file {fp}')
  print(f'>> nlen: {len(items)}')

  texts = [it['text'] for it in items]
  texts.sort()

  save_fp = OUT_PATH / 'testset-text.txt'
  print(f'>> save file: {save_fp}')
  with open(save_fp, 'w', encoding='utf-8') as fh:
    for t in texts:
      fh.write(t)
      fh.write('\n')


if __name__ == '__main__':
  stats_trainset()
  stats_testset()
